{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pnn4rDWGqDZL"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "l534d35Gp68G"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3TI3Q3XBesaS"
      },
      "source": [
        "# 训练检查点"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yw_a0iGucY8z"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://tensorflow.google.cn/guide/checkpoint\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\" class=\"\"> 在 TensorFlow.org 上查看</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/checkpoint.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\" class=\"\"> 在 Google Colab 中运行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/checkpoint.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\" class=\"\"> 在 GitHub 上查看源代码</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/guide/checkpoint.ipynb\" class=\"\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LeDp7dovcbus"
      },
      "source": [
        "“保存 TensorFlow 模型”这一短语通常表示保存以下两种元素之一：\n",
        "\n",
        "1. 检查点，或\n",
        "2. SavedModel。\n",
        "\n",
        "检查点可以捕获模型使用的所有参数（`tf.Variable` 对象）的确切值。检查点不包含对模型所定义计算的任何描述，因此通常仅在将使用保存参数值的源代码可用时才有用。\n",
        "\n",
        "另一方面，除了参数值（检查点）之外，SavedModel 格式还包括对模型所定义计算的序列化描述。这种格式的模型独立于创建模型的源代码。因此，它们适合通过 TensorFlow Serving、TensorFlow Lite、TensorFlow.js 或者使用其他编程语言（C、C++、Java、Go、Rust、C# 等 TensorFlow API）编写的程序进行部署。\n",
        "\n",
        "本文介绍用于编写和读取检查点的 API。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U0nm8k-6xfh2"
      },
      "source": [
        "## 设置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VEvpMYAKsC4z"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OEQCseyeC4Ev"
      },
      "outputs": [],
      "source": [
        "class Net(tf.keras.Model):\n",
        "  \"\"\"A simple linear model.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    super(Net, self).__init__()\n",
        "    self.l1 = tf.keras.layers.Dense(5)\n",
        "\n",
        "  def call(self, x):\n",
        "    return self.l1(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "utqeoDADC5ZR"
      },
      "outputs": [],
      "source": [
        "net = Net()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5vsq3-pffo1I"
      },
      "source": [
        "## 从 `tf.keras` 训练 API 保存\n",
        "\n",
        "请参阅 [`tf.keras` 保存和恢复指南](./keras/overview.ipynb#save_and_restore)。\n",
        "\n",
        "`tf.keras.Model.save_weights` 可以保存一个 TensorFlow 检查点。 "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SuhmrYPEl4D_"
      },
      "outputs": [],
      "source": [
        "net.save_weights('easy_checkpoint')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XseWX5jDg4lQ"
      },
      "source": [
        "## 编写检查点\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jpZPz76ZP3K"
      },
      "source": [
        "TensorFlow 模型的持久状态存储在 `tf.Variable` 对象中。这些对象可以直接构造，但通常会通过像 `tf.keras.layers` 或 `tf.keras.Model` 这样的高级 API 创建。\n",
        "\n",
        "管理变量的最简单方法是将它们附加到 Python 对象，然后引用这些对象。\n",
        "\n",
        "`tf.train.Checkpoint`、`tf.keras.layers.Layer` 和 `tf.keras.Model` 的子类会自动跟踪分配给其特性的变量。下面的示例构造了一个简单的线性模型，然后编写检查点，其中包含该模型所有变量的值。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x0vFBr_Im73_"
      },
      "source": [
        "您可以使用 `Model.save_weights` 轻松保存模型检查点。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHTJ1JzxCi8a"
      },
      "source": [
        "### 手动创建检查点"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6cF9fqYOCrEO"
      },
      "source": [
        "#### 设置"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fNjf9KaLdIRP"
      },
      "source": [
        "为了帮助演示 `tf.train.Checkpoint` 的所有功能， 下面定义了一个小数据集和优化步骤："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tSNyP4IJ9nkU"
      },
      "outputs": [],
      "source": [
        "def toy_dataset():\n",
        "  inputs = tf.range(10.)[:, None]\n",
        "  labels = inputs * 5. + tf.range(5.)[None, :]\n",
        "  return tf.data.Dataset.from_tensor_slices(\n",
        "    dict(x=inputs, y=labels)).repeat().batch(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ICm1cufh_JH8"
      },
      "outputs": [],
      "source": [
        "def train_step(net, example, optimizer):\n",
        "  \"\"\"Trains `net` on `example` using `optimizer`.\"\"\"\n",
        "  with tf.GradientTape() as tape:\n",
        "    output = net(example['x'])\n",
        "    loss = tf.reduce_mean(tf.abs(output - example['y']))\n",
        "  variables = net.trainable_variables\n",
        "  gradients = tape.gradient(loss, variables)\n",
        "  optimizer.apply_gradients(zip(gradients, variables))\n",
        "  return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vxzGpHRbOVO6"
      },
      "source": [
        "#### 创建检查点对象\n",
        "\n",
        "要手动创建检查点，您将需要 `tf.train.Checkpoint` 对象。您想要为对象设置检查点的位置将设置为此对象的特性。\n",
        "\n",
        "`tf.train.CheckpointManager` 也有助于管理多个检查点。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ou5qarOQOWYl"
      },
      "outputs": [],
      "source": [
        "opt = tf.keras.optimizers.Adam(0.1)\n",
        "dataset = toy_dataset()\n",
        "iterator = iter(dataset)\n",
        "ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)\n",
        "manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ZbYSD4uCy96"
      },
      "source": [
        "#### 训练模型并为模型设置检查点"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NP9IySmCeCkn"
      },
      "source": [
        "以下训练循环可创建模型和优化器的实例，然后将它们收集到 `tf.train.Checkpoint` 对象中。它在每批数据上循环调用训练步骤，并定期将检查点写入磁盘。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbCS5A6K1VSH"
      },
      "outputs": [],
      "source": [
        "def train_and_checkpoint(net, manager):\n",
        "  ckpt.restore(manager.latest_checkpoint)\n",
        "  if manager.latest_checkpoint:\n",
        "    print(\"Restored from {}\".format(manager.latest_checkpoint))\n",
        "  else:\n",
        "    print(\"Initializing from scratch.\")\n",
        "\n",
        "  for _ in range(50):\n",
        "    example = next(iterator)\n",
        "    loss = train_step(net, example, opt)\n",
        "    ckpt.step.assign_add(1)\n",
        "    if int(ckpt.step) % 10 == 0:\n",
        "      save_path = manager.save()\n",
        "      print(\"Saved checkpoint for step {}: {}\".format(int(ckpt.step), save_path))\n",
        "      print(\"loss {:1.2f}\".format(loss.numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ik3IBMTdPW41"
      },
      "outputs": [],
      "source": [
        "train_and_checkpoint(net, manager)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2wzcc1xYN-sH"
      },
      "source": [
        "#### 恢复和继续训练"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lw1QeyRBgsLE"
      },
      "source": [
        "在第一次设置检查点后，您可以传递新的模型和管理器，但需要从您离开的地方开始训练："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UjilkTOV2PBK"
      },
      "outputs": [],
      "source": [
        "opt = tf.keras.optimizers.Adam(0.1)\n",
        "net = Net()\n",
        "dataset = toy_dataset()\n",
        "iterator = iter(dataset)\n",
        "ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)\n",
        "manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)\n",
        "\n",
        "train_and_checkpoint(net, manager)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dxJT9vV-2PnZ"
      },
      "source": [
        "`tf.train.CheckpointManager` 对象会删除旧的检查点。上面配置为仅保留最近的三个检查点。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3zmM0a-F5XqC"
      },
      "outputs": [],
      "source": [
        "print(manager.checkpoints)  # List the three remaining checkpoints"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qwlYDyjemY4P"
      },
      "source": [
        "这些路径（如 `'./tf_ckpts/ckpt-10'`）不是磁盘上的文件，而是一个 `index` 文件和一个或多个包含变量值的数据文件的前缀。这些前缀被分组到一个单独的 `checkpoint` 文件 (`'./tf_ckpts/checkpoint'`) 中，其中 `CheckpointManager` 保存其状态。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t1feej9JntV_"
      },
      "outputs": [],
      "source": [
        "!ls ./tf_ckpts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DR2wQc9x6b3X"
      },
      "source": [
        "<a id=\"loading_mechanics\"></a>\n",
        "\n",
        "## 加载机制\n",
        "\n",
        "TensorFlow 通过从加载的对象开始遍历带命名边的有向计算图来将变量与检查点值匹配。边名称通常来自对象中的特性名称，例如 `self.l1 = tf.keras.layers.Dense(5)` 中的 `\"l1\"`。`tf.train.Checkpoint` 使用其关键字参数名称，如 `tf.train.Checkpoint(step=...)` 中的 `\"step\"`。\n",
        "\n",
        "上面示例中的依赖图如下所示：\n",
        "\n",
        "![Visualization of the dependency graph for the example training loop](https://tensorflow.google.cn/images/guide/whole_checkpoint.svg)\n",
        "\n",
        "优化器为红色，常规变量为蓝色，优化器插槽变量为橙色。其他节点（例如代表 `tf.train.Checkpoint`  的节点）为黑色。\n",
        "\n",
        "插槽变量是优化器状态的一部分，但是是为特定变量创建的。例如，上面的 `'m'` 边对应于动量，Adam 优化器会针对每个变量跟踪该动量。只有在同时保存变量和优化器时，才会将插槽变量保存到检查点中，并因此保存虚线边。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VpY5IuanUEQ0"
      },
      "source": [
        "在 `tf.train.Checkpoint` 对象上调用 `restore()` 会对请求的恢复进行排队，一旦有来自 `Checkpoint` 对象的匹配路径，就会立即恢复变量值。例如，通过网络和层重构一个指向上面所定义模型的路径后，我们便可以仅加载该模型中的偏差。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wmX2AuyH7TVt"
      },
      "outputs": [],
      "source": [
        "to_restore = tf.Variable(tf.zeros([5]))\n",
        "print(to_restore.numpy())  # All zeros\n",
        "fake_layer = tf.train.Checkpoint(bias=to_restore)\n",
        "fake_net = tf.train.Checkpoint(l1=fake_layer)\n",
        "new_root = tf.train.Checkpoint(net=fake_net)\n",
        "status = new_root.restore(tf.train.latest_checkpoint('./tf_ckpts/'))\n",
        "print(to_restore.numpy())  # We get the restored value now"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GqEW-_pJDAnE"
      },
      "source": [
        "这些新对象的依赖图是我们在上面编写的较大检查点的更小子图。它只包含偏差和 `tf.train.Checkpoint` 用来计算检查点数量的保存计数器。\n",
        "\n",
        "![Visualization of a subgraph for the bias variable](https://tensorflow.google.cn/images/guide/partial_checkpoint.svg)\n",
        "\n",
        "`restore()` 返回一个状态对象，该对象具有可选的断言。我们在新 `Checkpoint` 中创建的所有对象均已恢复，因此 `status.assert_existing_objects_matched()` 传递。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P9TQXl81Dq5r"
      },
      "outputs": [],
      "source": [
        "status.assert_existing_objects_matched()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GoMwf8CFDu9r"
      },
      "source": [
        "检查点中有许多尚未匹配的对象，包括层的内核和优化器的变量。`status.assert_consumed()` 仅在检查点和程序完全匹配时传递，并在此处引发异常。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KCcmJ-2j9RUP"
      },
      "source": [
        "### 延迟恢复\n",
        "\n",
        "当输入形状可用时，TensorFlow 中的 `Layer` 对象可能会将变量创建延迟到变量的首次调用。例如，`Dense` 层内核的形状取决于该层的输入和输出形状，因此，作为构造函数参数所需的输出形状没有足够的信息来单独创建变量。由于调用 `Layer` 还会读取变量的值，必须在变量的创建与其首次使用之间进行恢复。\n",
        "\n",
        "为支持这种习惯用法，`tf.train.Checkpoint` 会对尚不具有匹配变量的恢复进行排队。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TXYUCO3v-I72"
      },
      "outputs": [],
      "source": [
        "delayed_restore = tf.Variable(tf.zeros([1, 5]))\n",
        "print(delayed_restore.numpy())  # Not restored; still zeros\n",
        "fake_layer.kernel = delayed_restore\n",
        "print(delayed_restore.numpy())  # Restored"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-DWhJ3glyobN"
      },
      "source": [
        "### 手动检查检查点\n",
        "\n",
        "`tf.train.list_variables` 可以列出检查点键和检查点中变量的形状。检查点键是上面显示的计算图中的路径。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RlRsADTezoBD"
      },
      "outputs": [],
      "source": [
        "tf.train.list_variables(tf.train.latest_checkpoint('./tf_ckpts/'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fxk_BnZ4W1b"
      },
      "source": [
        "### 列表和字典跟踪\n",
        "\n",
        "对于像 `self.l1 = tf.keras.layers.Dense(5)` 一样的直接特性赋值，将列表和字典分配给特性会跟踪其内容。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rfaIbDtDHAr_"
      },
      "outputs": [],
      "source": [
        "save = tf.train.Checkpoint()\n",
        "save.listed = [tf.Variable(1.)]\n",
        "save.listed.append(tf.Variable(2.))\n",
        "save.mapped = {'one': save.listed[0]}\n",
        "save.mapped['two'] = save.listed[1]\n",
        "save_path = save.save('./tf_list_example')\n",
        "\n",
        "restore = tf.train.Checkpoint()\n",
        "v2 = tf.Variable(0.)\n",
        "assert 0. == v2.numpy()  # Not restored yet\n",
        "restore.mapped = {'two': v2}\n",
        "restore.restore(save_path)\n",
        "assert 2. == v2.numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UTKvbxHcI3T2"
      },
      "source": [
        "您可能会注意到列表和字典的包装器对象。这些包装器是可设置检查点版本的基础数据结构。就像基于特性的加载一样，这些包装器会在将变量添加到容器后立即恢复它的值。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s0Uq1Hv5JCmm"
      },
      "outputs": [],
      "source": [
        "restore.listed = []\n",
        "print(restore.listed)  # ListWrapper([])\n",
        "v1 = tf.Variable(0.)\n",
        "restore.listed.append(v1)  # Restores v1, from restore() in the previous cell\n",
        "assert 1. == v1.numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OxCIf2J6JyQ8"
      },
      "source": [
        "相同的跟踪会自动应用于 `tf.keras.Model` 的子类，并且可用于跟踪层列表等用途。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zGG1tOM0L6iM"
      },
      "source": [
        "## 使用 Estimator 保存基于对象的检查点\n",
        "\n",
        "请参阅 [ Estimator 指南](https://tensorflow.google.cn/guide/estimator)。\n",
        "\n",
        "默认情况下，Estimator 使用变量名而不是前面几部分中介绍的对象计算图来保存检查点。`tf.train.Checkpoint` 将接受基于名称的检查点，但是在将模型的一部分移到 Estimator 的 `model_fn` 外部时，变量名称可能会更改。保存基于对象的检查点可以更轻松地在 Estimator 内训练模型，然后在外部使用。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-8AMJeueNyoM"
      },
      "outputs": [],
      "source": [
        "import tensorflow.compat.v1 as tf_compat"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T6fQsBzJQN2y"
      },
      "outputs": [],
      "source": [
        "def model_fn(features, labels, mode):\n",
        "  net = Net()\n",
        "  opt = tf.keras.optimizers.Adam(0.1)\n",
        "  ckpt = tf.train.Checkpoint(step=tf_compat.train.get_global_step(),\n",
        "                             optimizer=opt, net=net)\n",
        "  with tf.GradientTape() as tape:\n",
        "    output = net(features['x'])\n",
        "    loss = tf.reduce_mean(tf.abs(output - features['y']))\n",
        "  variables = net.trainable_variables\n",
        "  gradients = tape.gradient(loss, variables)\n",
        "  return tf.estimator.EstimatorSpec(\n",
        "    mode,\n",
        "    loss=loss,\n",
        "    train_op=tf.group(opt.apply_gradients(zip(gradients, variables)),\n",
        "                      ckpt.step.assign_add(1)),\n",
        "    # Tell the Estimator to save \"ckpt\" in an object-based format.\n",
        "    scaffold=tf_compat.train.Scaffold(saver=ckpt))\n",
        "\n",
        "tf.keras.backend.clear_session()\n",
        "est = tf.estimator.Estimator(model_fn, './tf_estimator_example/')\n",
        "est.train(toy_dataset, steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tObYHnrrb_mL"
      },
      "source": [
        "随后，`tf.train.Checkpoint` 可以从其 `model_dir` 加载 Estimator 的检查点。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Q6IP3Y_wb-fs"
      },
      "outputs": [],
      "source": [
        "opt = tf.keras.optimizers.Adam(0.1)\n",
        "net = Net()\n",
        "ckpt = tf.train.Checkpoint(\n",
        "  step=tf.Variable(1, dtype=tf.int64), optimizer=opt, net=net)\n",
        "ckpt.restore(tf.train.latest_checkpoint('./tf_estimator_example/'))\n",
        "ckpt.step.numpy()  # From est.train(..., steps=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "knyUFMrJg8y4"
      },
      "source": [
        "## 总结\n",
        "\n",
        "TensorFlow 对象提供了一种简单的自动机制来保存和恢复它们所使用变量的值。\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "checkpoint.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
